from transformers import AutoModelForCausalLM, AutoTokenizer
import time
import chainlit as cl
from fastapi import FastAPI
from chainlit.utils import mount_chainlit
from chainlit.types import ThreadDict
from openai import AsyncOpenAI
from mcp import ClientSession
from typing import Dict, Optional
from fastapi import Request, Response
from chainlit.input_widget import Select, Switch, Slider
import pandas as pd
import plotly.graph_objects as go
import matplotlib.pyplot as plt
from datasets import load_dataset, Audio

 
@cl.set_starters
async def set_starters():
    '''初始化界面 提示'''
    return [
        cl.Starter(
            label="加载数据集",
            message="加载数据集",
            icon="/public/idea.svg",
        ),

        cl.Starter(
            label="Explain superconductors",
            message="Explain superconductors like I'm five years old.",
            icon="/public/learn.svg",
        ),
        cl.Starter(
            label="Python script for daily email reports",
            message="Write a script to automate sending daily email reports in Python, and walk me through how I would set it up.",
            icon="/public/terminal.svg",
            command="code",
        ),
        cl.Starter(
            label="Text inviting friend to wedding",
            message="Write a text asking a friend to be my plus-one at a wedding next month. I want to keep it super short and casual, and offer an out.",
            icon="/public/write.svg",
        )
    ]

def 加载数据集():
    print('start 加载数据集')
    dataset = load_dataset("cornell-movie-review-data/rotten_tomatoes")
    dataset=dataset["train"]
    dataset.rename_column("act", "labels")

    print(dataset)
    print(dataset[-1]["prompt"])
    

@cl.on_message
async def on_message(message: cl.Message):
   加载数据集()
   pass


